﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using TorchSharp.Modules;
using static TorchSharp.torch;

namespace NgpSharp
{
    public class NeRFNetwork : NeRFRender
    {
        protected int num_layers;
        protected int hidden_dim;
        protected int geo_feat_dim;
        //protected HashEmbedder embedder;
        protected GridEncoder encoder;
        protected int in_dim;
        protected ModuleList<Linear> sigma_net;

        protected int num_layers_color;
        protected int hidden_dim_color;
        protected SHEncoder encoder_dir;
        protected int in_dim_dir;
        protected ModuleList<Linear> color_net;

        public NeRFNetwork(string name, int num_layers = 2, int hidden_dim = 64, int geo_feat_dim = 15, int num_layers_color = 3, int hidden_dim_color = 64, 
            int num_layers_bg = 2, int hidden_dim_bg = 64, 
            int density_scale = 1, float min_near = 0.2F, float density_thresh = 0.01F) 
            : base(name, density_scale, min_near, density_thresh)
        {
            this.num_layers = num_layers;
            this.hidden_dim = hidden_dim;
            this.geo_feat_dim = geo_feat_dim;

            //HashEmbedder embedder = get_embedder();
            GridEncoder encoder = get_encoder(desired_resolution: 2048);
            this.encoder = encoder;
            this.in_dim = encoder.OutDim;

            List<Linear> sigma_net = [];
            for (int l = 0; l < num_layers; l++)
            {
                int in_dim;
                if (l == 0) in_dim = this.in_dim;
                else in_dim = hidden_dim;
                int out_dim;
                if (l == num_layers - 1) out_dim = 1 + geo_feat_dim;
                else out_dim = hidden_dim;
                sigma_net.Add(torch.nn.Linear(in_dim, out_dim, hasBias: false));
            }
            this.sigma_net = torch.nn.ModuleList([.. sigma_net]);

            this.num_layers_color = num_layers_color;
            this.hidden_dim_color = hidden_dim_color;
            this.encoder_dir = get_shencoder();
            this.in_dim_dir = encoder_dir.output_dim;

            List<Linear> color_net = [];
            for(int l = 0; l < num_layers_color; l++)
            {
                int in_dim;
                if (l == 0) in_dim = in_dim_dir + geo_feat_dim;
                else in_dim = hidden_dim_color;
                int out_dim;
                if (l == num_layers_color - 1) out_dim = 3;
                else out_dim = hidden_dim_color;
                color_net.Add(torch.nn.Linear(in_dim, out_dim, hasBias: false));
            }
            this.color_net = torch.nn.ModuleList([.. color_net]);
            RegisterComponents();
        }



        public override Tuple<Tensor, Tensor> forward(torch.Tensor x, torch.Tensor d, Tensor aabb)
        {
            //# x: [N, 3], in [-bound, bound]
            //# d: [N, 3], nomalized in [-1, 1]
            using var c = torch.NewDisposeScope();

            //# sigma
            x = encoder.call(x, aabb);

            torch.Tensor h = x;
            for(int l = 0; l < num_layers; l++)
            {
                h = sigma_net[l].call(h);
                if (l != num_layers - 1)
                    h = nn.functional.relu(h, inplace: true);
            }
            Tensor sigma = Activation.apply(h[TensorIndex.Ellipsis, 0]);
            Tensor geo_feat = h[TensorIndex.Ellipsis, TensorIndex.Slice(1..)];

            d = encoder_dir.call(d, 1);
            h = torch.cat([d, geo_feat], dim: -1);
            for(int l = 0; l < num_layers_color; l++)
            {
                h = color_net[l].call(h);
                if (l != num_layers_color - 1)
                    h = nn.functional.relu(h, inplace: true);
            }

            //# sigmoid activation for rgb
            Tensor color = torch.sigmoid(h);

            return new Tuple<Tensor, Tensor>(sigma, color);
        }

        public override Dictionary<string, Tensor> Density(Tensor x, Tensor aabb)
        {
            using var c = torch.NewDisposeScope();

            //# x: [N, 3], in [-bound, bound]
            x = encoder.call(x, aabb);
            Tensor h = x;

            for (int l = 0; l < num_layers; l++)
            {
                h = sigma_net[l].call(h);
                if (l != num_layers - 1)
                    h = nn.functional.relu(h, inplace: true);
            }

            Tensor simga = Activation.apply(h[TensorIndex.Ellipsis, 0]);
            Tensor geo_feat = h[TensorIndex.Ellipsis, TensorIndex.Slice(1)];

            return new Dictionary<string, Tensor>()
            {
                { "sigma", simga.MoveToOuterDisposeScope() },
                { "geo_feat", geo_feat.MoveToOuterDisposeScope() }
            };
        }

        public override Dictionary<string, Tensor> Color(Tensor x, Tensor d, Tensor? mask = null, Tensor? geo_feat = null)
        {
            using var c = torch.NewDisposeScope();

            Tensor? rgbs = null;
            if(mask is not null)
            {
                rgbs = torch.zeros(mask.shape[0], 3, dtype: x.dtype, device: x.device);
                //# in case of empty mask
                if(!Equals(mask.any(), torch.tensor(true).to(x.device)))
                {
                    return new Dictionary<string, Tensor>()
                    {
                        { "color", rgbs.MoveToOuterDisposeScope() }
                    };
                }
                x = x[mask];
                d = d[mask];
                geo_feat = geo_feat![mask];
            }

            d = encoder_dir.call(d, 1);
            Tensor h = torch.cat([d, geo_feat!], dim: -1);
            for (int l = 0; l < num_layers_color; l++)
            {
                h = color_net[l].call(h);
                if (l != num_layers_color - 1)
                    h = torch.nn.functional.relu(h, inplace: true);
            }

            //# sigmoid activation for rgb
            h = torch.sigmoid(h);

            if (mask is not null)
                rgbs![mask] = h.to(rgbs.dtype);
            else rgbs = h;
            return new Dictionary<string, Tensor>()
            {
                { "color", rgbs.MoveToOuterDisposeScope() }
            };
        }

        protected override nn.Module _to(Device device, ScalarType dtype, bool non_blocking)
        {
            To(device);
            return base._to(device, dtype, non_blocking);
        }

        protected override nn.Module _to(DeviceType deviceType, int deviceIndex, bool non_blocking)
        {
            Device device = torch.device(deviceIndex);
            To(device);
            return base._to(deviceType, deviceIndex, non_blocking);
        }

        /// <summary>
        /// 
        /// </summary>
        /// <returns></returns>
        public IEnumerable<Parameter> GetParameters()
        {
            List<Parameter> result = new List<Parameter>();
            result.AddRange(encoder.parameters());
            result.AddRange(sigma_net.parameters());
            result.AddRange(encoder_dir.parameters());
            result.AddRange(color_net.parameters());
            return result;
        }

        GridEncoder get_encoder(int input_dim = 3,
            int num_levels = 16, int level_dim = 2, 
            int base_resolution = 16, int log2_hashmap_size = 19, 
            int desired_resolution = 512)
        {
            var encoder = new GridEncoder("hashgrid", input_dim, num_levels,
                level_dim, base_resolution, desired_resolution, log2_hashmap_size);
            return encoder;
        }
        SHEncoder get_shencoder(int input_dim = 3, int degree = 4)
        {
            var encoder = new SHEncoder("sphere_harmonics", input_dim, degree);
            return encoder;
        }

        HashEmbedder get_embedder(int n_levels = 16, 
            int n_features_per_level = 2, int log2_hashmap_size = 19, 
            int base_resolution = 16, int finest_resolution = 512)
        {
            var encoder = new HashEmbedder("hashgrid", n_levels, n_features_per_level,
                log2_hashmap_size, base_resolution, finest_resolution);
            return encoder;
        }

        void To(Device device)
        {
            foreach(var item in sigma_net)
            {
                item.to(device);
            }
            foreach(var item in color_net)
            {
                item.to(device);
            }
        }
    }
}
